import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestRegressor

# 读取数据
df = pd.read_excel("enhanced_sales_data.xlsx")

# 按款号和时间排序
df.sort_values(['spu_code', 'days_since_launch'], inplace=True)


def train_and_predict(df, last_day, forecast_days=14):
    """滚动训练并预测未来14天销量"""
    # 训练数据：最近30天数据
    train_data = df[(df['days_since_launch'] > last_day - 30) &
                    (df['days_since_launch'] <= last_day)]

    if len(train_data) < 15:  # 数据量不足时使用简单平均
        avg_sale = train_data['sale_qty'].mean()
        return [avg_sale] * forecast_days

    # 特征工程
    X = train_data[['sum_cart_total', 'new_sale_qty', 'days_since_launch']]
    y = train_data['sale_qty']

    # 训练模型
    model = RandomForestRegressor(n_estimators=50, random_state=42)
    model.fit(X, y)

    # 预测未来14天
    predictions = []
    for day in range(1, forecast_days + 1):
        # 创建预测日的特征
        next_day = last_day + day
        # 使用最后一条记录的静态特征（加购数和新品销量不变）
        last_row = train_data.iloc[-1][['sum_cart_total', 'new_sale_qty']]
        features = pd.DataFrame({
            'sum_cart_total': [last_row['sum_cart_total']],
            'new_sale_qty': [last_row['new_sale_qty']],
            'days_since_launch': [next_day]
        })

        # 预测当日销量
        pred = model.predict(features)[0]
        predictions.append(max(0, pred))  # 确保非负

    return predictions


# 获取所有款号
all_spu = df['spu_code'].unique()
results = []

# 对每个款号进行预测
for spu in all_spu:
    spu_data = df[df['spu_code'] == spu]
    last_day = spu_data['days_since_launch'].max()

    # 预测未来14天销量
    pred_sales = train_and_predict(spu_data, last_day)

    # 保存结果
    for i, sale in enumerate(pred_sales, 1):
        results.append({
            'spu_code': spu,
            'days_since_launch': last_day + i,
            'predicted_sale_qty': sale
        })

# 转换为DataFrame并保存
results_df = pd.DataFrame(results)
results_df.to_excel("predicted_sales.xlsx", index=False)